import matplotlib.pyplot as plt
from pylab import *
from mpl_finance import candlestick_ochl
from matplotlib import gridspec, patches


def DrawBars(ax, poslist, valuelist, up="salmon", down="darkseagreen"):
    width = 0.7
    for i in range(len(poslist)):
        pos = poslist[i]
        value = valuelist[i]
        color = up
        if value < 0:
            color = down
        patch = ax.add_patch(
            patches.Rectangle(
                (pos - width / 2, 0),  # (x,y)
                width,  # width
                value,  # height
                edgecolor=color,
                facecolor=color
            ),
        )


def DrawBOLL(ax, datas):
    poslist = [one["pos"] for one in datas]
    upband = [one["bband"][0] for one in datas]
    ave = [one["bband"][1] for one in datas]
    dnband = [one["bband"][2] for one in datas]
    ax.plot(poslist, ave, color="b", alpha=0.5)
    ax.plot(poslist, upband, color='r', alpha=0.5)
    ax.plot(poslist, dnband, color="g", alpha=0.5)


def DrawMacd(ax, datas):
    poslist = [one["pos"] for one in datas]
    macd = [one["macd"][0] for one in datas]
    macdsignal = [one["macd"][1] for one in datas]
    macdhist = [one["macd"][2] * 2 for one in datas]
    DrawBars(ax, poslist, macdhist)
    ax.plot(poslist, macd, color="r")
    ax.plot(poslist, macdsignal, color="g")


def DrawSAR(ax, datas):
    poslist = [one["pos"] for one in datas]
    psar = [one["sar"] for one in datas]
    ax.plot(poslist, psar, "o", markersize=3, markeredgewidth=0, markeredgecolor='r', markerfacecolor='r')


def DrawSTOCH(ax, datas):
    poslist = [one["pos"] for one in datas]
    slowk = [one["stoch"][0] for one in datas]
    slowd = [one["stoch"][1] for one in datas]
    ax.plot(poslist, slowk, color="r")
    ax.plot(poslist, slowd, color="g")


def DrawTranPoint(ax, datas, pointoffet=5):
    for one in datas:
        if "tran_mark" in one:
            if one["tran_mark"] == 1:
                ax.annotate("%0.3f" % one["tran_mark"], xy=(one["pos"], one["Low"] - pointoffet),
                            xytext=(one["pos"], one["Low"] - pointoffet - 2),
                            arrowprops=dict(facecolor='red', width=1, shrink=0.25))
            if one["tran_mark"] == -1:
                ax.annotate("%0.3f" % one["tran_mark"], xy=(one["pos"], one["High"] + pointoffet),
                            xytext=(one["pos"], one["High"] + pointoffet + 2),
                            arrowprops=dict(facecolor='green', width=1, shrink=0.25))


def exDrawLine(axs, spandata):
    # DrawBOLL(axs[0],spandata)
    DrawSAR(axs[0], spandata)
    DrawMacd(axs[2], spandata)
    DrawSTOCH(axs[1], spandata)
    # DrawTranPoint(axs[0],spandata)


def Show(datalist, exDrawLine=None, pd='FXUSDJPY', qt_type=15):
    datas = []
    for i in range(len(datalist)):
        datalist[i]["pos"] = i
        datas.append([i, datalist[i]["Open"], datalist[i]["Close"], datalist[i]["High"], datalist[i]["Low"],
                      datalist[i]["Date"]])
    idtodate = {}
    idtodate = {one[0]: one[5] for one in datas}

    try:
        fig = plt.figure()
    except Exception as e:
        print(e)
        return
    fig.subplots_adjust(top=0.98, bottom=0.02, left=0.05, right=0.98)
    fig.subplots_adjust(hspace=0.01)
    gs = gridspec.GridSpec(3, 1, height_ratios=[4, 1, 1])
    ax = plt.subplot(gs[0])
    ax.set_title('Capacity Index')
    ax.set_ylabel("USD/oz")
    ax1 = plt.subplot(gs[1], sharex=ax)
    ax2 = plt.subplot(gs[2], sharex=ax)

    def getDate(x, position):
        if x in idtodate:
            t = idtodate[x]
            return t
        return ""

    ax.xaxis.set_major_formatter(FuncFormatter(getDate))
    bars = candlestick_ochl(ax, datas, width=0.6, colorup='r', colordown='g')

    if exDrawLine:
        exDrawLine((ax, ax1, ax2), datalist)

    ax.autoscale_view()
    ax1.autoscale_view()
    ax2.autoscale_view()

    ax.grid()
    ax1.grid()
    ax2.grid()
    ax.set_title(pd)
    ax.legend([pd], loc=0)

    mng = plt.get_current_fig_manager()
    # mng.window.showMaximized()
    plt.savefig("gan/%s - %s.png" % (pd, qt_type))
    # plt.pause(1)
    plt.close()


if __name__ == "__main__":
    datalist = fetchdata.loaddate("data/USDJPYdata.gz")
    Show(datalist[-500:], exDrawLine)
    # Show("data/goldmin.svr", "data/goldmin.data", True, 1200)
